package com.ambow.common.utils.mybatis;

import org.mybatis.generator.api.GeneratedXmlFile;
import org.mybatis.generator.api.IntrospectedColumn;
import org.mybatis.generator.api.IntrospectedTable;
import org.mybatis.generator.api.PluginAdapter;
import org.mybatis.generator.api.dom.java.Interface;
import org.mybatis.generator.api.dom.java.Method;
import org.mybatis.generator.api.dom.java.TopLevelClass;
import org.mybatis.generator.api.dom.xml.Attribute;
import org.mybatis.generator.api.dom.xml.Element;
import org.mybatis.generator.api.dom.xml.XmlElement;
import org.mybatis.generator.config.GeneratedKey;

import java.lang.reflect.Field;
import java.util.Iterator;
import java.util.List;

/**
 * Created by zhangyl on 2019/4/9
 */
public class MapperPlugin extends PluginAdapter {
    private static final String WARN = "当前文件为MybatisGenerator自动生成，重新生成时会被覆盖，请勿修改！（表结构变化时请重新生成）";
    @Override
    public boolean validate(List<String> list) {
        System.out.println("[INFO] MapperPlugin开始生成Mapper文件...");
        return true;
    }

    @Override
    public boolean sqlMapGenerated(GeneratedXmlFile sqlMap, IntrospectedTable introspectedTable) {
        try {
            Field field = sqlMap.getClass().getDeclaredField("isMergeable");
            field.setAccessible(true);
            field.setBoolean(sqlMap, false);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return true;
    }

    @Override
    public boolean modelBaseRecordClassGenerated(TopLevelClass topLevelClass, IntrospectedTable introspectedTable) {
        String tableName = introspectedTable.getAliasedFullyQualifiedTableNameAtRuntime();//数据库表名
        topLevelClass.addJavaDocLine("/**");
        topLevelClass.addJavaDocLine(" * " + WARN);
        topLevelClass.addJavaDocLine(" * " + "table:" + tableName);
        topLevelClass.addJavaDocLine(" */");

        //添加lombok注解
        topLevelClass.addImportedType("lombok.Getter");
        topLevelClass.addImportedType("lombok.Setter");
        topLevelClass.addAnnotation("@Getter");
        topLevelClass.addAnnotation("@Setter");
        return super.modelBaseRecordClassGenerated(topLevelClass, introspectedTable);
    }

    @Override
    public boolean modelGetterMethodGenerated(Method method, TopLevelClass topLevelClass, IntrospectedColumn introspectedColumn, IntrospectedTable introspectedTable, ModelClassType modelClassType) {
        return false;
    }

    @Override
    public boolean modelSetterMethodGenerated(Method method, TopLevelClass topLevelClass, IntrospectedColumn introspectedColumn, IntrospectedTable introspectedTable, ModelClassType modelClassType) {
        return false;
    }

    @Override
    public boolean clientInsertMethodGenerated(Method method, Interface interfaze, IntrospectedTable introspectedTable) {
        return false;
    }

    @Override
    public boolean sqlMapInsertElementGenerated(XmlElement element, IntrospectedTable introspectedTable) {
        return false;
    }

    @Override
    public boolean sqlMapInsertSelectiveElementGenerated(XmlElement element, IntrospectedTable introspectedTable) {
        replaceGeneratedKey(element, introspectedTable);
        return super.sqlMapInsertElementGenerated(element, introspectedTable);
    }

    /**
     * 将以下默认的主键返回形式替换为:keyProperty="id" useGeneratedKeys="true"
     * <selectKey keyProperty="id" order="AFTER" resultType="java.lang.Long">
     * SELECT LAST_INSERT_ID()
     * </selectKey>
     *
     * @param element
     * @param introspectedTable
     */
    private void replaceGeneratedKey(XmlElement element, IntrospectedTable introspectedTable) {
        GeneratedKey g = introspectedTable.getTableConfiguration().getGeneratedKey();
        if (g != null) {
            element.addAttribute(new Attribute("keyProperty", g.getColumn()));
            element.addAttribute(new Attribute("useGeneratedKeys", "true"));
            List<Element> els = element.getElements();
            Iterator<Element> its = els.iterator();
            while (its.hasNext()) {
                Element e = its.next();
                if (e instanceof XmlElement) {
                    XmlElement xe = (XmlElement) e;
                    if (xe.getName().equals("selectKey")) {
                        its.remove();
                    }
                }
            }
        }
    }
}
